from fairino import Robot
import time

class Robotmove:
    def __init__(self, ip_address='192.168.58.2'):
        self.position_dict = {
            "4": [11.331, -78.245, 108.114, -22.159, 170.936, 162.872],
            "1": [7.914, -71.742, 105.764, -16.765, 171.424, 162.874],
            "3": [-0.278, -74.616, 106.418, -10.217, 160.066, 162.873],
            "5": [4.303, -81.694, 110.953, -8.805, 159.192, 162.872],
            "6": [4.275, -95.327, 3.751, -8.896, 159.276, 162.871],  # 初始位置
            "2": [7.914, -71.742, 105.764, -16.765, 171.424, 162.874]
        }
        
        # 基础状态变量
        self.is_moving = False  # 运动状态
        self.is_connected = False  # 连接状态
        self.current_motion_error = None  # 运动错误状态
        self.last_target = None  # 最后目标位置
        self.last_motion_success = None  # 最后一次运动是否成功
        self.is_at_home = False  # 是否在初始位置（位置6）
        self.collision_strategy = None  # 碰撞策略状态
        self.collision_config_error = None  # 碰撞配置错误状态
        
        try:
            self.robot = Robot.RPC(ip_address)
            self.is_connected = True
            print("机械臂连接成功")
            
            # 初始化时设置碰撞策略为继续运动
            if self.set_crash_level(1):
                print("成功设置碰撞策略：碰撞后继续运动")
            else:
                print("设置碰撞策略失败")
                
        except Exception as e:
            print(f"无法连接到机械臂：{e}")
            self.is_connected = False
            self.current_motion_error = str(e)
        
        self.tool = 0
        self.user = 0

    def set_crash_level(self, block=0):
        """
        设置机械臂碰撞策略
        参数:
            block: 碰撞策略
                0: 碰撞后停止运动（默认）
                1: 碰撞后继续运动
        返回值:
            True: 设置成功
            False: 设置失败
        """
        try:
            # 设置碰撞策略
            error = self.robot.SetCollisionStrategy(strategy=block)
            
            if error != 0:
                print(f"设置碰撞策略失败，错误码: {error}")
                self.collision_strategy = None
                self.collision_config_error = error
                return False
            
            # 更新状态
            self.collision_strategy = block
            self.collision_config_error = None
            
            strategy_desc = "继续运动" if block == 1 else "停止运动"
            print(f"成功设置碰撞策略：碰撞后{strategy_desc}")
            return True
            
        except Exception as e:
            print(f"设置碰撞策略时发生错误: {str(e)}")
            self.collision_strategy = None
            self.collision_config_error = str(e)
            return False

    def get_collision_status(self):
        """
        获取当前碰撞策略状态
        返回值:
            dict: 包含碰撞策略相关的状态信息
        """
        return {
            "strategy": self.collision_strategy,
            "strategy_desc": "继续运动" if self.collision_strategy == 1 
                           else "停止运动" if self.collision_strategy == 0 
                           else "未设置",
            "config_error": self.collision_config_error,
            "is_configured": self.collision_strategy is not None
        }

    def move_joints(self, joint_id, velocity=30, blend_time=-1.0):
        """
        非阻塞关节运动
        参数:
            joint_id: 位置点ID
            velocity: 速度 (默认100)
            blend_time: 平滑时间 (默认0.5秒)
        返回值:
            True: 运动指令发送成功
            False: 运动指令发送失败
        """
        try:
            # 检查机械臂连接状态
            if not self.is_connected:
                print("机械臂未连接")
                self.current_motion_error = "机械臂未连接"
                return False

            # 检查当前运动状态
            if self.is_moving:
                print("机械臂正在运动中，请等待当前动作完成")
                return False

            # 检查位置是否有效
            if joint_id not in self.position_dict:
                print(f"错误：找不到关节位置 ID {joint_id}")
                self.current_motion_error = f"无效的位置ID: {joint_id}"
                return False

            # 发送运动指令
            error = self.robot.MoveJ(
                self.position_dict[joint_id], 
                self.tool, 
                self.user, 
                vel=velocity,
                blendT=blend_time  # 使用非阻塞运动
            )

            if error != 0:
                print(f"运动指令发送失败，错误码: {error}")
                self.current_motion_error = f"运动指令错误: {error}"
                self.last_motion_success = False
                return False
            
            # 更新状态
            self.is_moving = True
            self.last_target = joint_id
            self.current_motion_error = None
            self.last_motion_success = True
            
            # 更新初始位置状态（运动开始时先设为False）
            self.is_at_home = False
            
            print(f"成功发送运动指令到位置: {joint_id}")
            return True

        except Exception as e:
            error_msg = f"运动控制发生异常: {str(e)}"
            print(error_msg)
            self.current_motion_error = error_msg
            self.last_motion_success = False
            return False

    def update_motion_state(self):
        """
        更新运动状态（非阻塞）
        返回值:
            dict: 包含当前运动状态的信息
        """
        try:
            if not self.is_connected:
                return {
                    "is_moving": False,
                    "motion_done": False,
                    "error": "机械臂未连接",
                    "last_target": None,
                    "success": False,
                    "is_at_home": False
                }

            if not self.is_moving:
                return {
                    "is_moving": False,
                    "motion_done": True,
                    "error": self.current_motion_error,
                    "last_target": self.last_target,
                    "success": self.last_motion_success,
                    "is_at_home": self.is_at_home
                }

            # 获取运动完成状态
            _, ret = self.robot.GetRobotMotionDone()
            
            if ret == 1:  # 运动完成
                self.is_moving = False
                # 如果最后的目标是位置6，更新初始位置状态
                if self.last_target == "6":
                    self.is_at_home = True
                
                return {
                    "is_moving": False,
                    "motion_done": True,
                    "error": None,
                    "last_target": self.last_target,
                    "success": True,
                    "is_at_home": self.is_at_home
                }
            
            # 运动未完成
            return {
                "is_moving": True,
                "motion_done": False,
                "error": None,
                "last_target": self.last_target,
                "success": None,
                "is_at_home": self.is_at_home
            }

        except Exception as e:
            error_msg = f"状态更新发生错误: {str(e)}"
            print(error_msg)
            self.current_motion_error = error_msg
            self.is_moving = False
            return {
                "is_moving": False,
                "motion_done": False,
                "error": error_msg,
                "last_target": self.last_target,
                "success": False,
                "is_at_home": self.is_at_home
            }

    def get_motion_status(self):
        """
        获取当前运动状态信息
        返回值:
            dict: 包含所有运动相关的状态信息
        """
        return {
            "is_connected": self.is_connected,
            "is_moving": self.is_moving,
            "current_error": self.current_motion_error,
            "last_target": self.last_target,
            "last_success": self.last_motion_success,
            "is_at_home": self.is_at_home,
            "collision_strategy": self.collision_strategy
        }

def test_robot_sequence():
    """
    完整的机械臂测试序列
    测试内容：
    1. 初始化和碰撞策略设置
    2. 基础运动测试
    3. 返回初始位置测试
    4. 错误处理测试
    """
    print("开始机械臂测试序列...")
    robot = Robotmove()
    
    # 1. 检查初始化状态
    init_status = robot.get_motion_status()
    collision_status = robot.get_collision_status()
    print("\n1. 初始化状态检查:")
    print(f"连接状态: {init_status['is_connected']}")
    print(f"碰撞策略: {collision_status['strategy_desc']}")
    
    if not init_status['is_connected']:
        print("机械臂连接失败，测试终止")
        return
        
    # 2. 基础运动测试
    print("\n2. 开始基础运动测试:")
    
    # 测试位置1
    print("\n移动到位置1:")
    if robot.move_joints("1"):
        while True:
            status = robot.update_motion_state()
            if status["motion_done"]:
                print(f"位置1运动完成，成功: {status['success']}")
                break
            time.sleep(0.5)
    else:
        print("发送运动指令失败")
    
    time.sleep(1)  # 等待一秒
    
    # 测试位置2
    print("\n移动到位置2:")
    if robot.move_joints("2"):
        while True:
            status = robot.update_motion_state()
            if status["motion_done"]:
                print(f"位置2运动完成，成功: {status['success']}")
                break
            time.sleep(0.5)
    
    time.sleep(1)  # 等待一秒
    
    # 3. 返回初始位置测试
    print("\n3. 测试返回初始位置:")
    if robot.move_joints("6"):
        while True:
            status = robot.update_motion_state()
            if status["motion_done"]:
                final_status = robot.get_motion_status()
                print(f"运动完成，是否在初始位置: {final_status['is_at_home']}")
                break
            time.sleep(0.5)
    
    # 4. 错误处理测试
    print("\n4. 错误处理测试:")
    
    # 测试无效位置
    print("\n测试无效位置:")
    if not robot.move_joints("999"):
        print("无效位置测试通过：成功拒绝无效位置")
    
    # 测试运动状态检查
    print("\n测试运动状态检查:")
    if robot.move_joints("1"):
        # 立即尝试另一个运动
        if not robot.move_joints("2"):
            print("运动状态检查通过：成功阻止了重复运动")
        
        # 等待第一个运动完成
        while robot.update_motion_state()["is_moving"]:
            time.sleep(0.1)
    
    # 最终回到初始位置
    print("\n测试结束，返回初始位置:")
    if robot.move_joints("6"):
        while True:
            status = robot.update_motion_state()
            if status["motion_done"]:
                final_status = robot.get_motion_status()
                print(f"最终状态:")
                print(f"- 在初始位置: {final_status['is_at_home']}")
                print(f"- 运动状态: {'完成' if not final_status['is_moving'] else '运动中'}")
                print(f"- 错误状态: {final_status['current_error'] or '无错误'}")
                break
            time.sleep(0.1)

def test_collision_handling():
    """
    碰撞处理测试
    """
    print("\n开始碰撞处理测试...")
    robot = Robotmove()
    
    # 检查初始碰撞策略
    collision_status = robot.get_collision_status()
    print(f"初始碰撞策略: {collision_status['strategy_desc']}")
    
    # 尝试改变碰撞策略
    print("\n尝试更改碰撞策略到停止模式:")
    if robot.set_crash_level(0):
        collision_status = robot.get_collision_status()
        print(f"新的碰撞策略: {collision_status['strategy_desc']}")
    
    # 恢复碰撞策略
    print("\n恢复碰撞策略到继续运动模式:")
    if robot.set_crash_level(1):
        collision_status = robot.get_collision_status()
        print(f"最终碰撞策略: {collision_status['strategy_desc']}")

if __name__ == "__main__":
    print("=== 开始机械臂完整测试 ===")
    
    # 运行基础功能测试
    test_robot_sequence()
    
    # 运行碰撞处理测试
    test_collision_handling()
    
    print("\n=== 测试完成 ===")